import os
from PIL import Image, ImageChops
import numpy as np


class Results:
    @staticmethod
    def _ensure_dir(path: str):
        if path and not os.path.exists(path):
            os.makedirs(path, exist_ok=True)

    @staticmethod
    def _load_image(path: str) -> Image.Image:
        # Convert to RGB with white background to make comparison robust
        img = Image.open(path).convert("RGBA")
        bg = Image.new("RGBA", img.size, (255, 255, 255, 255))
        img = Image.alpha_composite(bg, img)
        return img.convert("RGB")

    @staticmethod
    def _similarity(a: Image.Image, b: Image.Image) -> float:
        a_arr = np.asarray(a, dtype=np.float32)
        b_arr = np.asarray(b, dtype=np.float32)
        if a_arr.shape != b_arr.shape:
            return 0.0
        mse = np.mean((a_arr - b_arr) ** 2)
        if mse <= 1e-12:
            return 1.0
        # Simple NMSE-based similarity: 1 - (MSE / 255^2)
        sim = max(0.0, 1.0 - mse / (255.0 ** 2))
        return float(sim)

    @staticmethod
    def _save_diff(actual_path: str,
                   reference_path: str,
                   name: str,
                   diff_dir: str,
                   diff_prefix: str = "diff_") -> float:
        """
        Save a visual diff image and return similarity ratio in [0, 1].
        If reference is missing, copy actual to reference and return 1.0
        to ease first-time bootstrap of baselines.
        """
        Results._ensure_dir(os.path.dirname(actual_path))
        Results._ensure_dir(os.path.dirname(reference_path))
        Results._ensure_dir(diff_dir)

        a = Results._load_image(actual_path)
        if not os.path.exists(reference_path):
            # First run: provision baseline automatically
            a.save(reference_path)
            # No diff needed; fully "similar" to itself.
            return 1.0

        r = Results._load_image(reference_path)

        # Pad to same canvas size if needed
        if a.size != r.size:
            W = max(a.width, r.width)
            H = max(a.height, r.height)
            na = Image.new("RGB", (W, H), (255, 255, 255))
            nr = Image.new("RGB", (W, H), (255, 255, 255))
            na.paste(a, (0, 0))
            nr.paste(r, (0, 0))
            a, r = na, nr

        sim = Results._similarity(a, r)

        diff = ImageChops.difference(a, r)
        # Amplify differences to be more visible
        amp = 4
        diff = Image.eval(diff, lambda p: min(255, p * amp))
        diff.save(os.path.join(diff_dir, f"{diff_prefix}{name}"))

        return sim
